"""
内容分析模块 - 专注于内容区域的提取和分析

此模块负责从二值化图像中提取内容区域，分析区域特征，
并提供内容区域的分组和可视化功能。
"""
import cv2
import numpy as np
from typing import List, Tuple, Optional
from dataclasses import dataclass


@dataclass
class ContentRegion:
    """内容区域数据类"""
    box: Tuple[int, int, int, int]  # 左上x, 左上y, 宽度, 高度
    area: int  # 区域面积（像素数）
    centroid: Tuple[int, int]  # 中心点坐标 (x, y)
    binary_mask: np.ndarray  # 二值化掩码


class ContentAnalyzer:
    """
    内容分析器类，专注于内容区域的提取和分析

    功能包括：
    - 从二值化图像中提取内容区域
    - 内容区域的特征分析
    - 内容区域的分组
    - 可视化功能
    """

    def __init__(self, min_content_area: int = 100):
        """
        初始化内容分析器

        参数:
            min_content_area: 最小内容区域面积（像素数）
        """
        self.min_content_area = min_content_area

    def extract_content_regions(self, binary_image: np.ndarray,
                                min_area: Optional[int] = None) -> List[ContentRegion]:
        """
        从二值化图像中提取内容区域

        参数:
            binary_image: 二值化图像
            min_area: 最小区域面积，如果为None则使用默认值

        返回:
            List[ContentRegion]: 内容区域列表
        """
        if binary_image is None or binary_image.size == 0:
            return []

        min_area = min_area if min_area is not None else self.min_content_area

        # 形态学操作：连接相近的文本区域
        kernel = np.ones((3, 3), np.uint8)
        dilated = cv2.dilate(binary_image, kernel, iterations=1)

        # 查找轮廓
        contours, _ = cv2.findContours(
            dilated,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )

        content_regions = []

        for contour in contours:
            # 计算轮廓面积
            area = int(cv2.contourArea(contour))

            # 过滤掉过小的区域
            if area < min_area:
                continue

            # 获取边界框
            x, y, w, h = cv2.boundingRect(contour)

            # 过滤掉太细长的区域
            aspect_ratio = w / h if h > 0 else 0
            if aspect_ratio > 10 or aspect_ratio < 0.1:
                continue

            # 计算质心
            centroid = self._calculate_centroid(contour, x, y, w, h)

            # 创建区域掩码
            mask = self._create_region_mask(binary_image, contour)

            # 创建内容区域对象
            region = ContentRegion(
                box=(x, y, w, h),
                area=area,
                centroid=centroid,
                binary_mask=mask
            )

            content_regions.append(region)

        return content_regions

    def _calculate_centroid(self, contour: np.ndarray, x: int, y: int,
                            w: int, h: int) -> Tuple[int, int]:
        """计算轮廓的质心"""
        M = cv2.moments(contour)
        if M["m00"] == 0:
            cX = x + w // 2
            cY = y + h // 2
        else:
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
        return (cX, cY)

    def _create_region_mask(self, binary_image: np.ndarray,
                            contour: np.ndarray) -> np.ndarray:
        """创建区域掩码"""
        mask = np.zeros_like(binary_image)
        cv2.drawContours(mask, [contour], 0, (255), -1)
        return mask

    def group_content_regions(self, content_regions: List[ContentRegion],
                              max_distance: int = 20,
                              include_images: bool = True) -> List[List[ContentRegion]]:
        """
        将内容区域分组（例如，将属于同一题目的文本和图像分在一组）

        参数:
            content_regions: 内容区域列表
            max_distance: 最大分组距离
            include_images: 是否包含图像区域

        返回:
            List[List[ContentRegion]]: 分组后的内容区域列表
        """
        if not content_regions:
            return []

        # 根据面积和宽高比对区域进行分类
        text_regions, image_regions = self._classify_regions(
            content_regions, include_images)

        # 按照y坐标排序文本区域
        text_regions.sort(key=lambda r: r.box[1])

        # 分组文本区域
        groups = self._group_text_regions(text_regions, max_distance)

        # 将图像分配到最近的文本组
        self._assign_images_to_groups(groups, image_regions, max_distance)

        return groups

    def _classify_regions(self, content_regions: List[ContentRegion],
                          include_images: bool) -> Tuple[List[ContentRegion], List[ContentRegion]]:
        """根据特征分类区域"""
        text_regions = []
        image_regions = []

        for region in content_regions:
            x, y, w, h = region.box
            aspect_ratio = w / h if h > 0 else 0

            # 简单启发式：宽高比大于4或小于0.25的可能是图像
            if aspect_ratio > 4 or aspect_ratio < 0.25 or region.area > 10000:
                if include_images:
                    image_regions.append(region)
            else:
                text_regions.append(region)

        return text_regions, image_regions

    def _group_text_regions(self, text_regions: List[ContentRegion],
                            max_distance: int) -> List[List[ContentRegion]]:
        """根据垂直距离分组文本区域"""
        groups = []
        if not text_regions:
            return groups

        current_group = [text_regions[0]]

        for i in range(1, len(text_regions)):
            prev_region = text_regions[i-1]
            curr_region = text_regions[i]

            prev_bottom = prev_region.box[1] + prev_region.box[3]
            curr_top = curr_region.box[1]

            if curr_top - prev_bottom <= max_distance:
                # 属于同一组
                current_group.append(curr_region)
            else:
                # 开始新的组
                if current_group:
                    groups.append(current_group)
                current_group = [curr_region]

        # 添加最后一组
        if current_group:
            groups.append(current_group)

        return groups

    def _assign_images_to_groups(self, groups: List[List[ContentRegion]],
                                 image_regions: List[ContentRegion],
                                 max_distance: int) -> None:
        """将图像分配到最近的文本组"""
        for img_region in image_regions:
            best_group_idx = -1
            min_distance = float('inf')

            for i, group in enumerate(groups):
                for text_region in group:
                    # 计算中心点距离
                    img_cx, img_cy = img_region.centroid
                    text_cx, text_cy = text_region.centroid

                    dist = ((img_cx - text_cx)**2 + (img_cy - text_cy)**2)**0.5

                    if dist < min_distance:
                        min_distance = dist
                        best_group_idx = i

            if best_group_idx >= 0 and min_distance < max_distance * 2:
                groups[best_group_idx].append(img_region)
            else:
                # 如果图像离所有文本都很远，则单独一组
                groups.append([img_region])

    def draw_content_regions(self, image: np.ndarray, content_regions: List[ContentRegion],
                             color: Tuple[int, int, int] = (0, 255, 0),
                             thickness: int = 2) -> np.ndarray:
        """
        在图像上绘制内容区域

        参数:
            image: 输入图像
            content_regions: 内容区域列表
            color: 边界框颜色
            thickness: 边界框线条粗细

        返回:
            np.ndarray: 绘制了内容区域的图像
        """
        result_image = image.copy()

        for i, region in enumerate(content_regions):
            x, y, w, h = region.box
            cv2.rectangle(result_image, (x, y), (x+w, y+h), color, thickness)

            # 标注区域序号
            cv2.putText(
                result_image,
                f"#{i+1}",
                (x, y-5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                thickness
            )

        return result_image
